from __future__ import with_statement

import unittest

import random
import time
import sys
import logging

import numpy
import scipy

import discrete_wlc
import refolding_common
import refolding
import test_util

logger = logging.getLogger("refolding_test")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)


gb1_spec = refolding.Spec(refolding_common.gb1_kf0, refolding_common.gb1_alpha)
gb1_kf0_slice = slice(gb1_spec.kf0 * 0.5, gb1_spec.kf0 * 1.5, 50)

def make_nonnegative(a):
    a[a < 0] = 0
    return a

def normalise(t):
    m = t.shape[0]
    result = numpy.empty((m, m), dtype=refolding_common.float_type)
    for j in range(m):
        result[:, j] = t[:, j] / numpy.sqrt((t[:, j] ** 2).sum())
    return result

def generate_movements(reunfolding_ramp, hover, delay, size):
    z = make_nonnegative(numpy.random.normal(loc=hover, scale=5e-9, size=size))
    t = make_nonnegative(numpy.random.normal(loc=delay, scale=5e-3, size=size))
    caught = numpy.max(numpy.array([
                numpy.ceil((reunfolding_ramp + z) / refolding_common.gb1_unfolded_length).astype(int),
                numpy.random.randint(15, size=size),
                ]), axis=0)
    unfolded = numpy.ceil(
        ((reunfolding_ramp + z) - caught * refolding_common.gb1_folded_length) / refolding_common.gb1_loop_length).astype(int)
    assert (unfolded > 0).all()
    unfolded_cl = caught * refolding_common.gb1_folded_length + unfolded * refolding_common.gb1_loop_length
    return [refolding.Movement(z[i], t[i], unfolded[i], unfolded_cl[i], caught[i]) for i in range(size)]


class RefoldingTest(unittest.TestCase):

    def test_simplification_is_noop_for_zero_length_loop(self):
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, 0)
        realistic_model = refolding.Model(
            model.spec, model.kb_t_over_p, refolding_common.gb1_loop_length)
        k = 0.15
        hover = refolding.Hover(
            k, refolding_common.gb1_8_unfolded_length,
            refolding_common.gb1_8_folded_length + refolding_common.gb1_loop_length * 1.5, 8)
        t = 10e-3  # 1 second .. whatever
        realistic_actuals = numpy.empty(hover.unfolded + 1, dtype=refolding_common.float_type)
        for i in range(hover.unfolded + 1):
            timed_hover = refolding.TimedHover(hover, t, i)
            actual = refolding.log_likelihood_from_timed_hover(model, timed_hover)
            expected = refolding.simplified_log_likelihood_from_timed_hovers(model, [timed_hover])
            numpy.testing.assert_almost_equal(actual, expected)
            realistic_actual = refolding.log_likelihood_from_timed_hover(realistic_model, timed_hover)
            logger.debug(
                "i=%d actual=%r expected=%r realistic_actual=%r",
                i, actual, expected, realistic_actual)
            self.assertFalse(numpy.isnan(realistic_actual))
            realistic_actuals[i] = realistic_actual
        total_probability = numpy.exp(realistic_actuals).sum()
        numpy.testing.assert_approx_equal(total_probability, 1)
        logger.debug("total_probability=%r", total_probability)

    def test_generate_time_to_first_folding(self):
        expected_kf = 4.0
        kf0 = 2.0
        f = 10.0
        model = refolding.Spec(kf0, -numpy.log(expected_kf / kf0) / f)
        unfolded = 10
        size = 10000
        t = model.generate_time_to_first_folding(unfolded, f, size)
        numpy.testing.assert_approx_equal(1 / (unfolded * expected_kf), t.mean(), 2)
        numpy.testing.assert_approx_equal(1 / (unfolded * expected_kf), t.std(), 1)

    def test_generate_folded_numbers(self):
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, 0)
        # Spring constant only matter when z and force are non-zero.
        k = 0.15
        hover = refolding.Hover(k, refolding_common.gb1_8_unfolded_length, 0, 8)
        size = 1000
        two_ms_folded = model.generate_folded_numbers(hover, 2e-3, size)
        two_ms_actual_fraction = sum(two_ms_folded) / float(8 * size)
        self.assertTrue(abs(two_ms_actual_fraction - 0.76321725) < 0.05, two_ms_actual_fraction)
        fifteen_ms_folded = model.generate_folded_numbers(hover, 15e-3, size)
        fifteen_ms_actual_fraction = sum(fifteen_ms_folded) / float(8 * size)
        self.assertTrue(
            abs(fifteen_ms_actual_fraction - 0.99998125) < 0.001, fifteen_ms_actual_fraction)
        # 12 pN corresponds to z/lc = 0.0651.
        hover = refolding.Hover(
            k, refolding_common.gb1_8_unfolded_length,
            0.0651 * refolding_common.gb1_8_unfolded_length, 8)
        ten_ms_12_pN_folded = model.generate_folded_numbers(hover, 10e-3, size)
        ten_ms_12_pN_actual_fraction = sum(ten_ms_12_pN_folded) / float(8 * size)
        self.assertTrue(
            abs(ten_ms_12_pN_actual_fraction - 0.980336) < 0.02, ten_ms_12_pN_actual_fraction)
        # A more realistic model taking into account the CL increase results in fewer folding events.
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, refolding_common.gb1_loop_length)
        loop_ten_ms_12_pN_folded = model.generate_folded_numbers(hover, 10e-3, size)
        loop_ten_ms_12_pN_actual_fraction = sum(loop_ten_ms_12_pN_folded) / float(8 * size)
        self.assertTrue(
            abs(loop_ten_ms_12_pN_actual_fraction - 0.8571) < 0.02, loop_ten_ms_12_pN_actual_fraction)
        
    def test_simplified_log_likelihood(self):
        size = 50
        # If delta_x was known, the average of a single time point could be enough to compute kf0
        # by matching its average. Let's aim for 50% folding prob for either hover distance, which
        # should also allow us to estimate delta_x.
        # 110 nm ramp size, 43 nm hovering, 20 ms waiting or 64.5 nm, 80 ms
        # 85 nm ramp, (60 nm, 120 ms [9.49 pN]), (36.3 nm, 20 ms [6.05 pN])
        # hover, delay = 36.3e-9, 20e-3  # 60e-9, 120e-3
        movements = generate_movements(reunfolding_ramp=85e-9, hover=36.3e-9, delay=20e-3, size=size)
        k = 0.15
        hovers = [refolding.Hover(
                k, movement.unfolded_cl, movement.z, movement.unfolded) for movement in movements]
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, 0)
        folded = numpy.concatenate(
            [model.generate_folded_numbers(hovers[i], movements[i].t, 1) for i in range(size)])
        caught_mean = numpy.mean([movement.caught for movement in movements])
        unfolded_mean = numpy.mean([movement.unfolded for movement in movements])
        fraction_folded = folded.mean() / unfolded_mean
        logger.info(
            "folded/unfolded=%f caught.mean()=%f unfolded.mean()=%f",
            fraction_folded, caught_mean, unfolded_mean)
        self.assertTrue(abs(fraction_folded - 0.5) < 0.2)
        timed_hovers = [
            refolding.TimedHover(hovers[i], movements[i].t, folded[i]) for i in range(size)]
        y, z = [], []
        for factor in numpy.arange(0.1, 2.0, 0.1):
            my = refolding.Model(
                refolding.Spec(refolding_common.gb1_kf0 * factor, refolding_common.gb1_alpha),
                refolding_common.kb_t_over_p, 0)
            y.append((refolding.simplified_log_likelihood_from_timed_hovers(my, timed_hovers), factor))
            mz = refolding.Model(
                refolding.Spec(refolding_common.gb1_kf0, refolding_common.gb1_alpha * factor),
                refolding_common.kb_t_over_p, 0)
            z.append((refolding.simplified_log_likelihood_from_timed_hovers(mz, timed_hovers), factor))
        max_y = max(y)
        self.assertTrue(0.65 < max_y[1] < 1.35, max_y)
        max_z = max(z)
        self.assertTrue(0.65 < max_z[1] < 1.35, max_z)

    def test_simplified_partial_log_likelihood_partial_kf0(self):
        unfolded, folded, t, forces = map(numpy.array, [[10, 20], [5, 15], [1e-3, 100e-3], [15e-12, 5e-12]])
        for alpha in numpy.arange(0.21e-9, 21e-9, 5e-10):
            exp_minus_alpha_forces = numpy.exp(-alpha * forces)
            for kf0 in numpy.arange(200, 2000, 100.0):
                actual = refolding.simplified_partial_log_likelihood_partial_kf0(
                    unfolded, folded, t, forces, kf0, alpha)
                sll_kf0 = lambda kf0 : refolding.simplified_log_likelihood(
                        unfolded, folded, t, kf0 * exp_minus_alpha_forces)
                expected = scipy.misc.derivative(sll_kf0, kf0)
                self.assertTrue(expected < 0)
                numpy.testing.assert_approx_equal(actual, expected, significant=6)
                #
                actual = refolding.simplified_partial_2_log_likelihood_partial_kf0_2(
                    unfolded, folded, t, forces, kf0, alpha)
                partial_sll_kf0 = lambda kf0 : refolding.simplified_partial_log_likelihood_partial_kf0(
                        unfolded, folded, t, forces, kf0, alpha)
                expected = scipy.misc.derivative(partial_sll_kf0, kf0)
                numpy.testing.assert_approx_equal(actual, expected, significant=5)
                #
                actual = refolding.simplified_partial_log_likelihood_partial_alpha(
                    unfolded, folded, t, forces, kf0, alpha)
                sll_alpha = lambda alpha : refolding.simplified_log_likelihood(
                        unfolded, folded, t, kf0 * numpy.exp(-alpha * forces))
                expected = scipy.misc.derivative(sll_alpha, alpha)
                numpy.testing.assert_approx_equal(actual, expected, significant=4)

    def test_simplified_estimate_kf0(self):
        size = 100
        movements = generate_movements(reunfolding_ramp=85e-9, hover=36.3e-9, delay=20e-3, size=size)
        k = 0.15
        hovers = [refolding.Hover(
                k, movement.unfolded_cl, movement.z, movement.unfolded) for movement in movements]
        model = refolding.Model(gb1_spec, refolding_common.kb_t_over_p, 0)
        folded = numpy.concatenate(
            [model.generate_folded_numbers(hovers[i], movements[i].t, 1) for i in range(size)])

        unfolded = numpy.array([movement.unfolded for movement in movements])
        t = numpy.array([movement.t for movement in movements])
        forces = refolding.compute_forces(model.kb_t_over_p, hovers)

        simplified_optimal_kf0, simplified_optimal_sll = refolding.simplified_estimate_kf0(
            unfolded, folded, t, forces, refolding_common.gb1_alpha)
        logger.info(
            "optimal_simplified_kf0=%f expected_simplified_kf0=%f",
            simplified_optimal_kf0, refolding_common.gb1_kf0)
        self.assertTrue(abs(simplified_optimal_kf0 - model.kf0()) < 150)

    def test_estimate_kf0(self):
        size = 100
        movements = generate_movements(reunfolding_ramp=85e-9, hover=36.3e-9, delay=20e-3, size=size)
        k = 0.15
        hovers = [refolding.Hover(
                k, movement.unfolded_cl, movement.z, movement.unfolded) for movement in movements]
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, refolding_common.gb1_loop_length)
        folded = numpy.concatenate(
            [model.generate_folded_numbers(hovers[i], movements[i].t, 1) for i in range(size)])
        timed_hovers = [refolding.TimedHover(hovers[i], movements[i].t, folded[i]) for i in range(size)]

        simplified_optimal_kf0, _ = refolding.simplified_estimate_kf0_from_timed_hovers(
            model.kb_t_over_p, model.loop_length, timed_hovers, model.alpha())
        optimal_kf0, _ = refolding.estimate_kf0_from_timed_hovers(
            model.kb_t_over_p, model.loop_length, timed_hovers, model.alpha(),
            gb1_kf0_slice)
        logger.info(
            "simplified_optimal_kf0=%r optimal_kf0=%r expected_kf0=%r",
            simplified_optimal_kf0, optimal_kf0, refolding_common.gb1_kf0)
        self.assertTrue(abs(optimal_kf0 - model.kf0()) < 200)
        self.assertTrue(simplified_optimal_kf0 < optimal_kf0)

    def _test_simplified_estimate(self):
        fct = lambda kb_t_over_p, loop_length, timed_hovers, alpha_slice, kf0_slice: refolding.simplified_estimate_from_timed_hovers(
            kb_t_over_p, refolding_common.gb1_loop_length, timed_hovers)
        self.__estimate_and_check(0, fct)

    def test_estimate(self):
        self.__estimate_and_check(
            refolding_common.gb1_loop_length, refolding.estimate_from_timed_hovers)

    def __estimate_and_check(self, loop_length, optimisation_function):
        # Using the real refolding_common.gb1_loop_length instead of 0, delta_x changes by 0.2 nm only.
        model = refolding.Model(
            gb1_spec, refolding_common.kb_t_over_p, loop_length)
        timed_hovers = []
        folded = numpy.array([], dtype=int)
        # (85e-9, 36.3e-9, 20e-3), (85e-9, 60e-9, 120e-3)
        # Dimeri: (190e-9, 90e-9, 300e-3), (235e-9, 45e-9, 15e-3)
        # Possibility: (205e-9, 35e-9, 20e-3), (170e-9, 70e-9, 250e-3)
        for (reunfolding_ramp, hover, delay) in [(205e-9, 35e-9, 15e-3), (170e-9, 70e-9, 250e-3)]:
            some_size=100
            some_movements = generate_movements(
                reunfolding_ramp=reunfolding_ramp, hover=hover, delay=delay, size=some_size)
            k = 0.15
            some_hovers = [refolding.Hover(
                    k, movement.unfolded_cl, movement.z, movement.unfolded)
                           for movement in some_movements]
            some_folded = numpy.concatenate(
                [model.generate_folded_numbers(some_hovers[i], some_movements[i].t, 1)
                 for i in range(some_size)])

            for i in range(some_size):
                m = some_movements[i]
                extra_unfolded = 1  # -m.unfolded / 2
                u = m.unfolded + extra_unfolded
                u  = max(u, some_folded[i])
                timed_hovers.append(refolding.TimedHover(
                        refolding.Hover(k, m.unfolded_cl, m.z, u), m.t, some_folded[i]))
        logger.info(
            "folded_fraction=%f mean_unfolded=%f mean_folded=%f",
                    float(sum(timed_hover.folded for timed_hover in timed_hovers)) / sum(
                timed_hover.unfolded() for timed_hover in timed_hovers),
            numpy.mean([timed_hover.unfolded() for timed_hover in timed_hovers]),
            numpy.mean([timed_hover.folded for timed_hover in timed_hovers]))
        optimal_alpha, optimal_kf0, optimal_sll, (_, _) = optimisation_function(
            model.kb_t_over_p, model.loop_length, timed_hovers,
            slice(gb1_spec.alpha * 0.5, gb1_spec.alpha * 1.5, 5e-10 / refolding_common.kb_t),
            gb1_kf0_slice
            )
        logger.info(
            "optimal_alpha=%f (optimal_x=%f nm) optimal_kf0=%f optimal_sll=%f",
            optimal_alpha, optimal_alpha * refolding_common.kb_t * 1e9, optimal_kf0, optimal_sll)
        delta_x = abs(optimal_alpha - refolding_common.gb1_alpha) * refolding_common.kb_t
        self.assertTrue(delta_x < 0.4e-9, delta_x)
        self.assertTrue(
            abs(optimal_kf0 - refolding_common.gb1_kf0) < refolding_common.gb1_kf0,
            optimal_kf0)
        
    def test_multiple_exp_integral_single_and_double_integral(self):
        t = 0.5
        exponent = numpy.array([0.2], dtype=refolding_common.float_type)
        actual = refolding.multiple_exp_integral(exponent, t)
        expected_0 = (numpy.exp(exponent[0] * t) - 1) / exponent[0]
        self.assertFalse(numpy.isnan(actual))
        numpy.testing.assert_approx_equal(actual, expected_0)
        numpy.testing.assert_approx_equal(
            refolding.slow_multiple_exp_integral(exponent, 0, t), expected_0)
        exponent = numpy.array([0.2, 0.3], dtype=refolding_common.float_type)
        expected_1 = 1.0 / exponent[1] * (
            numpy.exp(exponent[1] * t) * expected_0 - (numpy.exp(exponent.sum() * t) - 1) / exponent.sum())
        self.assertFalse(numpy.isnan(actual))
        actual = refolding.multiple_exp_integral(exponent, t)
        numpy.testing.assert_approx_equal(actual, expected_1)
        numpy.testing.assert_approx_equal(
            refolding.slow_multiple_exp_integral(exponent, 0, t), expected_1)

    def test_multiple_exp_integral_multiple_integral(self):
        t = 0.1
        # Could add more exponents but that takes computational time.
        c = [-32.082727772386392, -22.042934457254113, -14.028562142948449, -8.0624367986690562, -4.0279654245224128, -1.6477685668021049, -0.49972305415173918, -0.094121332999350993, -0.0078081662905916402]
        c = [0.05, 0.1, 0.07]
        exponent = numpy.array(c, dtype=refolding_common.float_type)
        actual = refolding.multiple_exp_integral(exponent, t)
        expected = refolding.slow_multiple_exp_integral(exponent, 0, t)
        numpy.testing.assert_approx_equal(actual, expected)


if __name__ == "__main__":
    numpy.seterr(divide="raise", invalid="raise")
    test_util.initialise_seeds()
    suite = unittest.TestLoader().loadTestsFromTestCase(RefoldingTest)
    unittest.TextTestRunner(verbosity=2).run(suite)
